import torch
from torch import Tensor
from modelscope import AutoTokenizer, AutoModel
import pandas as pd
import numpy as np
from tqdm import tqdm

from .settings import device, embedding_modelname

# 定义Qwen模型的最大序列长度
qwen_max_length = 8192

# 初始化tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained(embedding_modelname, padding_side='left')
model = AutoModel.from_pretrained(embedding_modelname)

# 将模型移动到指定设备并设置为评估模式
model = model.to(device)
model.eval()


def last_token_pool(last_hidden_states: Tensor,
                    attention_mask: Tensor) -> Tensor:
    """
    使用最后一token的隐藏状态作为整个序列的表示
    适用于处理不同长度的序列，考虑左填充情况

    Args:
        last_hidden_states (Tensor): 模型输出的最后隐藏状态
        attention_mask (Tensor): 注意力掩码，标识哪些位置是有效内容

    Returns:
        Tensor: 池化后的序列表示向量
    """
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]


def generate_all_embeddings(text_list, batch_size=4):
    """
    为文本列表中的所有文本生成嵌入向量
    使用批处理提高效率，并实时释放内存避免溢出

    Args:
        text_list (list): 需要生成嵌入向量的文本列表
        batch_size (int): 批处理大小，默认为4

    Returns:
        numpy.ndarray: 所有文本的嵌入向量矩阵
    """
    all_embeddings = []

    # 使用进度条显示处理进度
    for i in tqdm(range(0, len(text_list), batch_size)):
        batch_texts = text_list[i: i + batch_size]

        # 对文本进行分词和编码
        batch_dict = tokenizer(
            batch_texts,
            padding=True,
            truncation=True,
            max_length=qwen_max_length,
            return_tensors="pt",
        ).to(device)

        # 在不计算梯度的情况下生成嵌入向量
        with torch.no_grad():
            outputs = model(**batch_dict).last_hidden_state
            embeddings = last_token_pool(outputs, batch_dict['attention_mask'])

        # 将结果转移到CPU并转换为numpy数组
        embeddings_cpu = embeddings.cpu().numpy()
        all_embeddings.append(embeddings_cpu)

        # 释放GPU内存
        del batch_dict, outputs, embeddings
        torch.cuda.empty_cache()

    # 合并所有批次的嵌入向量
    final_embeddings = np.vstack(all_embeddings)
    return final_embeddings

#   是用于推理加速和节省内存的装饰器。
@torch.no_grad()
def max_min_diverse_subset(text_list, embeddings, k=7000):
    """
    贪心算法
    使用最大最小多样性算法选择文本子集
    确保选择的文本既具有代表性又具有多样性

    Args:
        text_list (list): 原始文本列表
        embeddings (numpy.ndarray): 所有文本的嵌入向量
        k (int): 需要选择的文本数量，默认为7000

    Returns:
        list: 被选中的文本索引列表
    """
    embeddings_tensor = torch.tensor(embeddings, device=device)
    # 对嵌入向量进行归一化处理
    embeddings_normalized = embeddings_tensor / torch.norm(embeddings_tensor, dim=1, keepdim=True)

    N = embeddings_normalized.shape[0]
    selected_indices = []

    # 随机选择一个以句号结尾的文本作为起始点
    while True:
        first_index = torch.randint(0, N, (1,)).item()
        if text_list[first_index].strip().endswith("。"):
            break
    selected_indices.append(first_index)

    # 计算所有文本与已选文本的最小余弦距离
    dist_to_selected = 1 - torch.mv(embeddings_normalized, embeddings_normalized[first_index])

    # 使用进度条显示选择进度
    pbar = tqdm(total=k - 1)
    while len(selected_indices) < k:
        # 选择距离当前已选集合最远的文本
        next_index = torch.argmax(dist_to_selected).item()

        # 跳过不以句号结尾的文本
        if not text_list[next_index].strip().endswith("。"):
            dist_to_selected[next_index] = 0
            continue

        selected_indices.append(next_index)

        # 更新所有文本与已选集合的最小距离
        dist_new = 1 - torch.mv(embeddings_normalized, embeddings_normalized[next_index])
        dist_to_selected = torch.minimum(dist_to_selected, dist_new)

        pbar.update(1)
    pbar.close()

    return selected_indices


# 读取法律文本数据
df = pd.read_csv("../../data/Laws_All.csv")
text_list = df["内容"].tolist()

# 为所有法律文本生成嵌入向量
embeddings = generate_all_embeddings(text_list)
# 保存嵌入向量到文件
np.save("../../data/Laws_Embeddings.npy", embeddings)

# 使用最大最小多样性算法选择文本子集
indices = max_min_diverse_subset(text_list, embeddings, 10000)
# 保存选中的文本到新文件
df_selected = df.iloc[indices]
df_selected.to_csv("../../data/Laws_Selected.csv", index=False, encoding="utf-8-sig")